#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF文件读取器
支持PDF文本提取、页面信息获取、图片提取等功能
"""

import os
import sys
import argparse
from pathlib import Path
import json

try:
    import PyPDF2
    PYPDF2_AVAILABLE = True
except ImportError:
    PYPDF2_AVAILABLE = False

try:
    import pdfplumber
    PDFPLUMBER_AVAILABLE = True
except ImportError:
    PDFPLUMBER_AVAILABLE = False

try:
    import fitz  # PyMuPDF
    PYMUPDF_AVAILABLE = True
except ImportError:
    PYMUPDF_AVAILABLE = False


class PDFReader:
    """PDF文件读取器类"""
    
    def __init__(self):
        self.supported_formats = ['.pdf']
        self.available_engines = []
        
        if PYPDF2_AVAILABLE:
            self.available_engines.append('PyPDF2')
        if PDFPLUMBER_AVAILABLE:
            self.available_engines.append('pdfplumber')
        if PYMUPDF_AVAILABLE:
            self.available_engines.append('PyMuPDF')
            
        if not self.available_engines:
            raise ImportError("未安装任何PDF处理库！请安装 PyPDF2、pdfplumber 或 PyMuPDF")
    
    def get_pdf_info(self, pdf_path):
        """
        获取PDF文件基本信息
        
        参数:
            pdf_path (str): PDF文件路径
        
        返回:
            dict: PDF文件信息
        """
        try:
            if not os.path.exists(pdf_path):
                raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")
            
            info = {
                'file_path': pdf_path,
                'file_size': os.path.getsize(pdf_path),
                'pages': 0,
                'title': '',
                'author': '',
                'subject': '',
                'creator': '',
                'producer': '',
                'creation_date': '',
                'modification_date': ''
            }
            
            # 优先使用PyMuPDF获取信息（功能最强）
            if PYMUPDF_AVAILABLE:
                doc = fitz.open(pdf_path)
                info['pages'] = doc.page_count
                metadata = doc.metadata
                info.update({
                    'title': metadata.get('title', ''),
                    'author': metadata.get('author', ''),
                    'subject': metadata.get('subject', ''),
                    'creator': metadata.get('creator', ''),
                    'producer': metadata.get('producer', ''),
                    'creation_date': metadata.get('creationDate', ''),
                    'modification_date': metadata.get('modDate', '')
                })
                doc.close()
            elif PYPDF2_AVAILABLE:
                with open(pdf_path, 'rb') as file:
                    pdf_reader = PyPDF2.PdfReader(file)
                    info['pages'] = len(pdf_reader.pages)
                    if pdf_reader.metadata:
                        info.update({
                            'title': pdf_reader.metadata.get('/Title', ''),
                            'author': pdf_reader.metadata.get('/Author', ''),
                            'subject': pdf_reader.metadata.get('/Subject', ''),
                            'creator': pdf_reader.metadata.get('/Creator', ''),
                            'producer': pdf_reader.metadata.get('/Producer', ''),
                            'creation_date': str(pdf_reader.metadata.get('/CreationDate', '')),
                            'modification_date': str(pdf_reader.metadata.get('/ModDate', ''))
                        })
            elif PDFPLUMBER_AVAILABLE:
                with pdfplumber.open(pdf_path) as pdf:
                    info['pages'] = len(pdf.pages)
            
            return info
            
        except Exception as e:
            print(f"获取PDF信息失败: {str(e)}")
            return None
    
    def extract_text_pypdf2(self, pdf_path, start_page=0, end_page=None):
        """使用PyPDF2提取文本"""
        text_content = []
        
        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            total_pages = len(pdf_reader.pages)
            
            if end_page is None:
                end_page = total_pages
            else:
                end_page = min(end_page, total_pages)
            
            for page_num in range(start_page, end_page):
                page = pdf_reader.pages[page_num]
                text = page.extract_text()
                text_content.append({
                    'page': page_num + 1,
                    'text': text
                })
                
        return text_content
    
    def extract_text_pdfplumber(self, pdf_path, start_page=0, end_page=None):
        """使用pdfplumber提取文本（支持表格）"""
        text_content = []
        
        with pdfplumber.open(pdf_path) as pdf:
            total_pages = len(pdf.pages)
            
            if end_page is None:
                end_page = total_pages
            else:
                end_page = min(end_page, total_pages)
            
            for page_num in range(start_page, end_page):
                page = pdf.pages[page_num]
                text = page.extract_text()
                
                # 尝试提取表格
                tables = page.extract_tables()
                
                page_content = {
                    'page': page_num + 1,
                    'text': text if text else '',
                    'tables': tables if tables else []
                }
                text_content.append(page_content)
                
        return text_content
    
    def extract_text_pymupdf(self, pdf_path, start_page=0, end_page=None):
        """使用PyMuPDF提取文本（最强功能）"""
        text_content = []
        
        doc = fitz.open(pdf_path)
        total_pages = doc.page_count
        
        if end_page is None:
            end_page = total_pages
        else:
            end_page = min(end_page, total_pages)
        
        for page_num in range(start_page, end_page):
            page = doc[page_num]
            text = page.get_text()
            
            # 获取页面详细信息
            page_info = {
                'page': page_num + 1,
                'text': text,
                'width': page.rect.width,
                'height': page.rect.height,
                'images': len(page.get_images()),
                'links': len(page.get_links())
            }
            text_content.append(page_info)
        
        doc.close()
        return text_content
    
    def extract_text(self, pdf_path, engine='auto', start_page=0, end_page=None):
        """
        提取PDF文本内容
        
        参数:
            pdf_path (str): PDF文件路径
            engine (str): 使用的引擎 ('auto', 'PyPDF2', 'pdfplumber', 'PyMuPDF')
            start_page (int): 起始页码（从0开始）
            end_page (int): 结束页码（不包含）
        
        返回:
            list: 页面文本内容列表
        """
        try:
            if not os.path.exists(pdf_path):
                raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")
            
            print(f"正在提取PDF文本: {pdf_path}")
            print(f"可用引擎: {self.available_engines}")
            
            # 自动选择最佳引擎
            if engine == 'auto':
                if PYMUPDF_AVAILABLE:
                    engine = 'PyMuPDF'
                elif PDFPLUMBER_AVAILABLE:
                    engine = 'pdfplumber'
                elif PYPDF2_AVAILABLE:
                    engine = 'PyPDF2'
                else:
                    raise ImportError("没有可用的PDF处理引擎")
            
            print(f"使用引擎: {engine}")
            
            # 根据引擎提取文本
            if engine == 'PyMuPDF' and PYMUPDF_AVAILABLE:
                return self.extract_text_pymupdf(pdf_path, start_page, end_page)
            elif engine == 'pdfplumber' and PDFPLUMBER_AVAILABLE:
                return self.extract_text_pdfplumber(pdf_path, start_page, end_page)
            elif engine == 'PyPDF2' and PYPDF2_AVAILABLE:
                return self.extract_text_pypdf2(pdf_path, start_page, end_page)
            else:
                raise ValueError(f"引擎 {engine} 不可用或未安装")
                
        except Exception as e:
            print(f"文本提取失败: {str(e)}")
            return []
    
    def extract_images(self, pdf_path, output_dir="extracted_images"):
        """
        提取PDF中的图片（需要PyMuPDF）
        
        参数:
            pdf_path (str): PDF文件路径
            output_dir (str): 图片输出目录
        
        返回:
            list: 提取的图片文件路径列表
        """
        if not PYMUPDF_AVAILABLE:
            print("图片提取需要PyMuPDF库，请安装: pip install PyMuPDF")
            return []
        
        try:
            if not os.path.exists(pdf_path):
                raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")
            
            # 创建输出目录
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            doc = fitz.open(pdf_path)
            image_files = []
            
            print(f"开始提取图片到目录: {output_dir}")
            
            for page_num in range(doc.page_count):
                page = doc[page_num]
                image_list = page.get_images()
                
                for img_index, img in enumerate(image_list):
                    xref = img[0]
                    pix = fitz.Pixmap(doc, xref)
                    
                    if pix.n - pix.alpha < 4:  # 只处理RGB或灰度图像
                        img_filename = f"page_{page_num + 1}_img_{img_index + 1}.png"
                        img_path = os.path.join(output_dir, img_filename)
                        pix.save(img_path)
                        image_files.append(img_path)
                        print(f"提取图片: {img_path}")
                    
                    pix = None
            
            doc.close()
            print(f"图片提取完成，共提取 {len(image_files)} 张图片")
            return image_files
            
        except Exception as e:
            print(f"图片提取失败: {str(e)}")
            return []
    
    def save_text_to_file(self, text_content, output_file, format='txt'):
        """
        将提取的文本保存到文件
        
        参数:
            text_content (list): 文本内容列表
            output_file (str): 输出文件路径
            format (str): 输出格式 ('txt', 'json')
        """
        try:
            if format == 'json':
                with open(output_file, 'w', encoding='utf-8') as f:
                    json.dump(text_content, f, ensure_ascii=False, indent=2)
            else:  # txt格式
                with open(output_file, 'w', encoding='utf-8') as f:
                    for page_content in text_content:
                        f.write(f"\n{'='*50}\n")
                        f.write(f"第 {page_content['page']} 页\n")
                        f.write(f"{'='*50}\n\n")
                        f.write(page_content['text'])
                        f.write('\n\n')
                        
                        # 如果有表格信息，也保存
                        if 'tables' in page_content and page_content['tables']:
                            f.write(f"\n--- 表格信息 ---\n")
                            for i, table in enumerate(page_content['tables']):
                                f.write(f"表格 {i+1}:\n")
                                for row in table:
                                    f.write('\t'.join([str(cell) if cell else '' for cell in row]) + '\n')
                                f.write('\n')
            
            print(f"文本已保存到: {output_file}")
            
        except Exception as e:
            print(f"保存文件失败: {str(e)}")


def main():
    """主函数，支持命令行调用"""
    parser = argparse.ArgumentParser(description='PDF文件读取器')
    parser.add_argument('pdf_file', help='要读取的PDF文件路径')
    parser.add_argument('-o', '--output', help='输出文件路径')
    parser.add_argument('-e', '--engine', default='auto', 
                      choices=['auto', 'PyPDF2', 'pdfplumber', 'PyMuPDF'],
                      help='使用的PDF处理引擎')
    parser.add_argument('-s', '--start-page', type=int, default=0, help='起始页码（从0开始）')
    parser.add_argument('-n', '--end-page', type=int, help='结束页码（不包含）')
    parser.add_argument('-f', '--format', default='txt', choices=['txt', 'json'],
                      help='输出格式')
    parser.add_argument('-i', '--info', action='store_true', help='仅显示PDF文件信息')
    parser.add_argument('--extract-images', action='store_true', help='提取PDF中的图片')
    parser.add_argument('--image-dir', default='extracted_images', help='图片输出目录')
    
    args = parser.parse_args()
    
    try:
        reader = PDFReader()
        
        # 显示PDF信息
        if args.info:
            info = reader.get_pdf_info(args.pdf_file)
            if info:
                print(f"\nPDF文件信息:")
                print(f"文件路径: {info['file_path']}")
                print(f"文件大小: {info['file_size']:,} 字节")
                print(f"页数: {info['pages']}")
                print(f"标题: {info['title']}")
                print(f"作者: {info['author']}")
                print(f"主题: {info['subject']}")
                print(f"创建者: {info['creator']}")
                print(f"生产者: {info['producer']}")
                print(f"创建时间: {info['creation_date']}")
                print(f"修改时间: {info['modification_date']}")
            return
        
        # 提取图片
        if args.extract_images:
            image_files = reader.extract_images(args.pdf_file, args.image_dir)
            if image_files:
                print(f"\n提取的图片文件:")
                for img_file in image_files:
                    print(f"  - {img_file}")
            return
        
        # 提取文本
        text_content = reader.extract_text(
            args.pdf_file,
            args.engine,
            args.start_page,
            args.end_page
        )
        
        if text_content:
            # 保存到文件
            if args.output:
                reader.save_text_to_file(text_content, args.output, args.format)
            else:
                # 显示前几页内容预览
                print(f"\n提取成功，共 {len(text_content)} 页")
                for i, page_content in enumerate(text_content[:3]):  # 只显示前3页
                    print(f"\n--- 第 {page_content['page']} 页 ---")
                    text = page_content['text'][:500]  # 只显示前500字符
                    print(text)
                    if len(page_content['text']) > 500:
                        print("...")
                
                if len(text_content) > 3:
                    print(f"\n... (共 {len(text_content)} 页，使用 -o 参数保存完整内容)")
        
    except Exception as e:
        print(f"程序运行出错: {str(e)}")


if __name__ == "__main__":
    # 如果直接运行脚本，提供交互式界面
    if len(sys.argv) == 1:
        print("=== PDF文件读取器 ===\n")
        
        try:
            reader = PDFReader()
            print(f"可用的PDF处理引擎: {reader.available_engines}")
            
            # 获取当前目录下的PDF文件
            pdf_files = list(Path('.').glob('*.pdf'))
            
            if pdf_files:
                print(f"\n发现以下PDF文件:")
                for i, file in enumerate(pdf_files, 1):
                    print(f"{i}. {file}")
                
                try:
                    choice = int(input(f"\n请选择要读取的文件 (1-{len(pdf_files)}): ")) - 1
                    if 0 <= choice < len(pdf_files):
                        selected_file = str(pdf_files[choice])
                        
                        # 显示文件信息
                        info = reader.get_pdf_info(selected_file)
                        if info:
                            print(f"\n文件信息:")
                            print(f"页数: {info['pages']}")
                            print(f"文件大小: {info['file_size']:,} 字节")
                            print(f"标题: {info['title']}")
                            print(f"作者: {info['author']}")
                        
                        # 选择操作
                        print(f"\n选择操作:")
                        print("1. 提取文本")
                        print("2. 提取图片") 
                        print("3. 显示详细信息")
                        
                        operation = input("请选择操作 (1-3): ").strip()
                        
                        if operation == '1':
                            # 提取文本
                            text_content = reader.extract_text(selected_file)
                            if text_content:
                                output_file = f"{Path(selected_file).stem}_extracted.txt"
                                reader.save_text_to_file(text_content, output_file)
                        
                        elif operation == '2':
                            # 提取图片
                            reader.extract_images(selected_file)
                        
                        elif operation == '3':
                            # 已经显示了基本信息，不需要额外操作
                            pass
                        
                        else:
                            print("无效的选择")
                    else:
                        print("无效的选择")
                        
                except (ValueError, KeyboardInterrupt):
                    print("操作已取消")
            else:
                print("当前目录下未找到PDF文件")
                pdf_file = input("请输入PDF文件路径: ").strip()
                if pdf_file and os.path.exists(pdf_file):
                    info = reader.get_pdf_info(pdf_file)
                    if info:
                        print(f"文件页数: {info['pages']}")
                        text_content = reader.extract_text(pdf_file)
                        if text_content:
                            output_file = f"{Path(pdf_file).stem}_extracted.txt"
                            reader.save_text_to_file(text_content, output_file)
                else:
                    print("文件不存在")
                    
        except ImportError as e:
            print(f"导入错误: {str(e)}")
            print("请安装必要的PDF处理库:")
            print("pip install PyPDF2 pdfplumber PyMuPDF")
        
        except Exception as e:
            print(f"程序运行出错: {str(e)}")
    else:
        main() 